package com.jboltai.event.ai.aiknowledge;

import cn.hutool.core.util.StrUtil;
import com.jboltai.capability.Templates;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.message.MessageFinishReason;
import com.jboltai.event.CompleteHandler;
import com.jboltai.event.EventType;
import com.jboltai.event.FailHandler;
import com.jboltai.event.SuccessHandler;
import com.jboltai.event.ai.aitext.AITextEvent;
import com.jboltai.event.handler.AIKnowledgeEventHandler;
import com.jboltai.resource.ai.AIModel;

import java.util.List;
import java.util.regex.Matcher;
import java.util.stream.Collectors;

/**
 * ai知识库聊天事件
 */
public class AIKnowledgeEvent extends AITextEvent<AIKnowledgeEvent, AIMessage, AIKnowledgeEventHandler> {

    /**
     * 无匹配时的默认回复
     */
    private static final String NO_MATCH_DEFAULT_REPLY = "抱歉，我无法回答您的问题";

    /**
     * 是否成功匹配知识库片段
     */
    private boolean matchSuccess;

    public AIKnowledgeEvent() {
        super(EventType.AI_KNOWLEDGE);

    }

    public AIKnowledgeEvent(AIModel model) {
        super(EventType.AI_KNOWLEDGE);
        this.aiChatEvent.setModel(model);
    }


    /**
     * 创建一个ai事件
     * @param model 使用哪个模型，具体的模型名称根据该枚举值的defaultModel获取
     * @param successHandler
     * @param failHandler
     */
    public AIKnowledgeEvent(AIModel model, SuccessHandler<AIKnowledgeEvent, AIMessage> successHandler, FailHandler failHandler) {
        super(EventType.AI_KNOWLEDGE);
        this.aiChatEvent.setModel(model);
        this.successHandler = successHandler;
        this.failHandler = failHandler;
    }
    /**
     * 创建一个ai事件
     * @param model 使用哪个模型，具体的模型名称根据该枚举值的defaultModel获取
     * @param successHandler
     * @param failHandler
     */
    public AIKnowledgeEvent(AIModel model, SuccessHandler<AIKnowledgeEvent, AIMessage> successHandler, FailHandler failHandler, CompleteHandler completeHandler) {
        super(EventType.AI_KNOWLEDGE);
        this.aiChatEvent.setModel(model);
        this.successHandler = successHandler;
        this.failHandler = failHandler;
        this.completeHandler = completeHandler;
    }



    /**
     * 提供的知识库文档
     */
    private List<String> background;
    /**
     * 是否仅知识库模式
     */
    private boolean knowledgeOnly;

    /**
     * 设置prompt模版
     */
    private String promptTpl;

    /**
     * 无匹配知识时的响应回答
     */
    private String noMatchReply;

    /**
     * 无匹配结果时的回调
     */
    private NoMatchHandler noMatchHandler;



    public String generateFinalPrompt() {
        String tpl = this.promptTpl;
        if (tpl == null) {
            //未指定模版
            tpl = this.knowledgeOnly ? Templates.CHAT_WITH_KNOWLEDGE_ONLY : Templates.CHAT_WITH_KNOWLEDGE;
        }

        tpl = tpl.replaceAll("\\{\\s*background\\s*}", Matcher.quoteReplacement((background == null || background.isEmpty()) ? "无" : background.stream().collect(Collectors.joining("\n"))) )
                .replaceAll("\\{\\s*question\\s*}", Matcher.quoteReplacement(this.oriPrompt))
                .replaceAll("\\{\\s*nomatch\\s*}", Matcher.quoteReplacement(getNoMatchReply()));
        return tpl;



    }


    public AIKnowledgeEvent onNoMatch(NoMatchHandler handler) {
        this.noMatchHandler = handler;
        return this;
    }

    public NoMatchHandler getNoMatchHandler() {
        return noMatchHandler;
    }

    /**
     * 获取提问
     * @return
     */
    public String getQuestion() {
        return this.oriPrompt;
    }

    /**
     * 获取背景知识
     * @return
     */
    public List<String> getBackground() {
        return background;
    }

    public AIKnowledgeEvent setBackground(List<String> background) {
        this.background = background;
        return this;
    }

    /**
     * 是否仅知识库模式
     * @return
     */
    public boolean getIsKnowledgeOnly() {
        return knowledgeOnly;
    }

    public AIKnowledgeEvent setKnowledgeOnly(boolean knowledgeOnly) {
        this.knowledgeOnly = knowledgeOnly;
        return this;
    }

    /**
     * 获取知识库模版
     * @return
     */
    public String getPromptTpl() {
        return promptTpl;
    }

    public AIKnowledgeEvent setPromptTpl(String promptTpl) {
        this.promptTpl = promptTpl;
        return this;
    }

    /**
     * 获取无响应时的回答
     * @return
     */
    public String getNoMatchReply() {
        return StrUtil.isBlank(this.noMatchReply) ? NO_MATCH_DEFAULT_REPLY : this.noMatchReply;
    }

    public AIKnowledgeEvent setNoMatchReply(String noMatchReply) {
        this.noMatchReply = noMatchReply;
        return this;
    }

    /**
     * 获取是否是stream请求，如果使用了functioncall ,会强制为false
     * @return
     */
    public boolean getIsStream() {
        return this.aiChatEvent.getIsStream();
    }

    public AIKnowledgeEvent setStream(boolean stream) {
        this.aiChatEvent.setStream(stream);
        return this;
    }



    @Override
    public String preCheck() {
        if (StrUtil.isBlank(this.oriPrompt) ) {
            return "未设置提问";
        }
        return null;
    }



    public MessageFinishReason getFinishReason() {
        return this.aiChatEvent.getFinishReason();
    }

    public Integer getPromptTokenCount() {
        return this.aiChatEvent.getPromptTokenCount();
    }

    public Integer getCompletionTokenCount() {
        return this.aiChatEvent.getCompletionTokenCount();
    }

    public Integer getTotalTokenCount() {
        return this.aiChatEvent.getTotalTokenCount();
    }

    public boolean isMatchSuccess() {
        return matchSuccess;
    }

    public void setMatchSuccess(boolean matchSuccess) {
        this.matchSuccess = matchSuccess;
    }
}
